#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/19 21:29
@Author  : thezehui@gmail.com
@File    : wikipedia_search.py.py
"""
from langchain_community.tools.wikipedia.tool import WikipediaQueryInput, WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from internal.lib.helper import add_attribute

class WikipediaArgsShema(BaseModel):
    """Input for the WikipediaQuery tool."""
    query: str = Field(description="query to look up on wikipedia")


@add_attribute("args_schema", WikipediaArgsShema)
def wikipedia_search(**kwargs) -> BaseTool:
    """返回维基百科搜索工具"""
    return WikipediaQueryRun(
        api_wrapper=WikipediaAPIWrapper(),
    )
